package com.study.spark.scala.dataframe

import java.lang.Long

import org.apache.orc.impl.TreeReaderFactory.LongTreeReader
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._
import org.apache.spark.sql.{Dataset, Row, SparkSession}

/**
  * UDAF
  * @author stephen
  * @create 2019-03-17 13:35
  * @since 1.0.0
  */
object UDAFDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("UDAF Demo")
      .master("local[2]")
      .getOrCreate()

    // 生成测试数据
    val range: Dataset[Long] = spark.range(1, 11)

    val geoMean = new GeoMean()

    // 方式一：sql
    //spark.udf.register("gm",geoMean)
    //range.createOrReplaceTempView("range")
    //val result = spark.sql("SELECT gm(id) result FROM range")

    // 方式二：DSL
    import spark.implicits._
    val result = range.groupBy().agg(geoMean($"id")).as("result")

    result.show()

    spark.stop()

  }
}


class GeoMean extends UserDefinedAggregateFunction {
  // 输入数据类型
  override def inputSchema: StructType = {
    StructType(List(StructField("value", DoubleType)))
  }

  // 产生中间结果的数据类型
  override def bufferSchema: StructType = {
    StructType(List(
      // 相乘之后返回的积
      StructField("product", DoubleType),
      // 参与运算数字的个数
      StructField("counts", LongType)
    ))

  }

  // 最终结果数据而理性
  override def dataType: DataType = DoubleType

  // 确保一致性
  override def deterministic: Boolean = true

  // 指定初始化值
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    // 相乘初始值
    buffer(0) = 1.0
    // 参与运算数字个数的初始值
    buffer(1) = 0L
  }

  // 每有一条数据参与运算就更新中间结果
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    // 每有一个数字参与运算就进行相乘
    buffer(0) = buffer.getDouble(0) * input.getDouble(0)
    // 更新参与运算数据个数
    buffer(1) = buffer.getLong(1) + 1L
  }

  // 全局聚合
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    // 每个分区计算结果进行相乘
    buffer1(0) = buffer1.getDouble(0) * buffer2.getDouble(0)
    // 每个分区参与计算的中间结果进行相加
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  // 计算最终结果
  override def evaluate(buffer: Row): Any = {
    math.pow(buffer.getDouble(0), 1.toDouble / buffer.getLong(1))
  }
}